
#include "pywinpcapmodule.h"

extern PyTypeObject pywinpcap_device_type;
extern PyTypeObject pywinpcap_packet_type;
extern PyObject * PyExc_PyWinPCapError;
extern PyObject * PyExc_PyWinPCapTimeoutError;

static void
pywinpcap_capture_dealloc(pywinpcap_capture_object * self)
{
	/* check arguments validity */
	if (NULL == self)
	{
		return ;
	}

	/* close the pcap object */
	if (NULL != self->m_pcap)
	{
		(void)pcap_close(self->m_pcap);
		self->m_pcap = NULL;
	}

	Py_XDECREF(self->m_device);
    self->ob_type->tp_free((PyObject*)self);
}

static PyObject * pywinpcap_capture_new(PyTypeObject * type,
										PyObject *args,
										PyObject *kwds)
{
    pywinpcap_capture_object * self = NULL;

    self = (pywinpcap_capture_object *)type->tp_alloc(type, 0);
    if (NULL == self)
	{
		return NULL;
	}

	self->m_device = NULL;
    self->m_pcap = NULL;

    return (PyObject *)self;
}

static int pywinpcap_capture_init(	pywinpcap_capture_object * self,
									PyObject * args,
									PyObject * kwds)
{
	char err_buff[PCAP_ERRBUF_SIZE];
	pywinpcap_device_object * device = NULL;
	int snaplen = 0xFFFF; /* max packet size */
	PyObject * promisc = NULL; /* promisc mode */
	int read_timeout = 1000;

	if (!PyArg_ParseTuple(args, "O|iOi", &device, &snaplen, &promisc, &read_timeout))
	{
		return -1;
	}
	
	/* set promisc default value */
	if (NULL == promisc)
	{
		promisc = Py_True;
	}

	/* ensure args type */
	if (!PyWinPcapDevice_Check(device) || !PyBool_Check(promisc))
	{
		return -1;
	}

	Py_BEGIN_ALLOW_THREADS

	/* open pcap reader object */
	self->m_pcap = pcap_open_live(	PyString_AS_STRING(device->m_device),
									snaplen,
									((PyBoolObject *)promisc)->ob_ival,
									read_timeout,
									err_buff);

	Py_END_ALLOW_THREADS

	if (NULL == self->m_pcap)
	{
		PyErr_SetString(PyExc_PyWinPCapError, err_buff);
		return -1;
	}
	
	Py_INCREF(device);
	self->m_device = device;

    return 0;
}

static PyObject * pywinpcap_capture_str(pywinpcap_capture_object * self)
{
	return PyString_FromFormat("<Capture '%s'>", PyString_AS_STRING(self->m_device->m_device));
}

static PyObject *
pywinpcap_capture_next(pywinpcap_capture_object * self, PyObject * args)
{
	PyStringObject * packet_data_str = NULL;
	struct pcap_pkthdr * header = NULL;
	const u_char * data = NULL;
	int pcap_ret = 0;

	if (!PyArg_ParseTuple(args, ""))
	{
		return NULL;
	}
	
	Py_BEGIN_ALLOW_THREADS

	pcap_ret = pcap_next_ex(self->m_pcap, &header, &data);

	Py_END_ALLOW_THREADS

	switch (pcap_ret)
	{
	case -2:
		PyErr_SetString(PyExc_PyWinPCapError, "EOF reached while waiting for packets");
		return NULL;
	case -1:
		PyErr_SetString(PyExc_PyWinPCapError, "Error waiting for packets");
		return NULL;
	case 0:
		PyErr_SetString(PyExc_PyWinPCapTimeoutError, "Timeout occured while waiting for packets");
		return NULL;
	case 1:
		break;
	default:
		PyErr_SetString(PyExc_PyWinPCapError, "Invalid error while waiting for packets");
		return NULL;
	}

	packet_data_str = (PyStringObject *)PyString_FromStringAndSize((char *)data, header->len);
	if (NULL == packet_data_str)
	{
		return NULL;
	}

	return PyObject_CallFunction(	(PyObject*)&pywinpcap_packet_type,
									"(Oii)",
									packet_data_str, 
									header->ts.tv_sec,
									header->ts.tv_usec);
}

PyDoc_STRVAR(pywinpcap_capture_next_doc,
"next() -> Get next packet available on device reader");

typedef
struct pkt_handler_context_rec
{
	pywinpcap_capture_object *	m_self;
	PyFunctionObject *			m_callback;
	PyObject *					m_context;
	int							m_pkt_received;
	BOOL						m_exception_occurred;
} pkt_handler_context_t;

void pkt_handler(	u_char * user,
					const struct pcap_pkthdr * pkt_header,
					const u_char * pkt_data)
{
	PyObject * packet_obj = NULL;
	PyObject * callback_ret = NULL;
	PyStringObject * packet_data_str = NULL;
	pkt_handler_context_t * context = (pkt_handler_context_t *)user;
	PyGILState_STATE gil_state;

	++context->m_pkt_received;

	gil_state = PyGILState_Ensure();

	packet_data_str = (PyStringObject *)PyString_FromStringAndSize(	(char *)pkt_data,
																	pkt_header->len);
	if (NULL == packet_data_str)
	{
		goto break_loop_with_exception;
	}

	/* create a packet object for the packet data */
	packet_obj = PyObject_CallFunction(	(PyObject *)&pywinpcap_packet_type,
										"(Oii)",
										packet_data_str,
										pkt_header->ts.tv_sec,
										pkt_header->ts.tv_usec);
	if (NULL == packet_obj)
	{
		goto break_loop_with_exception;
	}

	/* call the user callback */
	callback_ret = PyObject_CallFunction(	(PyObject*)context->m_callback,
											"(OOO)",
											context->m_self, 
											packet_obj,
											context->m_context);
	
	Py_XDECREF(packet_obj);

	if (NULL == callback_ret)
	{
		goto break_loop_with_exception;
	}

	/* check return value */
	if ((Py_None != callback_ret) && !PyBool_Check(callback_ret))
	{
		Py_XDECREF(callback_ret);
		
		PyErr_SetString(PyExc_PyWinPCapError, "User callback must return boolean type");

		goto break_loop_with_exception;
	}

	if (Py_False == callback_ret)
	{
		/* break out of loop */
		(void)pcap_breakloop(context->m_self->m_pcap);
	}

	Py_XDECREF(callback_ret);

	PyGILState_Release(gil_state);

	return;

break_loop_with_exception:
	PyGILState_Release(gil_state);

	/* break out of loop */
	(void)pcap_breakloop(context->m_self->m_pcap);

	context->m_exception_occurred = TRUE;

	return ;
}

static PyObject *
pywinpcap_capture_loop(pywinpcap_capture_object * self, PyObject * args)
{
	pkt_handler_context_t handler_context;
	BOOL decref_none = FALSE;
	int pkt_count = -1;
	int pcap_ret = 0;

	memset(&handler_context, 0, sizeof(pkt_handler_context_t));

	handler_context.m_exception_occurred = FALSE;

	/* parse arguments */
	if (!PyArg_ParseTuple(args, "O|Oi", &handler_context.m_callback, &handler_context.m_context, &pkt_count))
	{
		return NULL;
	}

	/* ensure the user callback is callable */
	if (!PyCallable_Check((PyObject *)handler_context.m_callback))
	{
		PyErr_SetString(PyExc_TypeError, "callback must be callable");
		return NULL;
	}

	/* set user context default value if needed */
	if (NULL == handler_context.m_context)
	{
		decref_none = TRUE;
		Py_INCREF(Py_None);
		handler_context.m_context = Py_None;
	}

	Py_INCREF(self);
	handler_context.m_self = self;

	Py_BEGIN_ALLOW_THREADS

	pcap_ret = pcap_loop(self->m_pcap, pkt_count, (pcap_handler)pkt_handler, (u_char *)&handler_context);

	Py_END_ALLOW_THREADS

	Py_DECREF(self);

	switch (pcap_ret)
	{
	case -2:
		/* if no exception was set, set default one */
		if (!handler_context.m_exception_occurred)
		{
			PyErr_SetString(PyExc_PyWinPCapError, "User broke out of loop");
		}

		return NULL;
	case -1:
		PyErr_SetString(PyExc_PyWinPCapError, "Error while looping on packets");
		return NULL;
	case 0: /* count is reached */
		break;
	default:
		PyErr_SetString(PyExc_PyWinPCapError, "Invalid error while waiting for packets");
		return NULL;
	}

	/* return how many packets we received */
	return PyInt_FromLong(handler_context.m_pkt_received);
}

PyDoc_STRVAR(pywinpcap_capture_loop_doc,
"loop(callback, context=None, cnt=-1) -> Loop over packets from the device and call the callback(capture_obj, packet_obj, context)");

static PyObject *
pywinpcap_capture_send(pywinpcap_capture_object * self, PyObject * args)
{
	PyObject * packet_data = NULL;
	int pcap_ret = 0;

	if (!PyArg_ParseTuple(args, "O", &packet_data))
	{
		return NULL;
	}

	if (!PyString_Check(packet_data))
	{
		PyErr_SetString(PyExc_Exception, "Packet data must be str");
		return NULL;
	}

	Py_BEGIN_ALLOW_THREADS

	pcap_ret = pcap_sendpacket(	self->m_pcap,
								(const u_char *)PyString_AS_STRING(packet_data),
								PyString_Size(packet_data));

	Py_END_ALLOW_THREADS

	/* send the packet */
	if (0 != pcap_ret)
	{
		PyErr_SetString(PyExc_Exception, "Error sending the packet");
		return NULL;
	}

	Py_INCREF(Py_None);
	return Py_None;
}

PyDoc_STRVAR(pywinpcap_capture_send_doc,
"send() -> Send a packet over the interface");

static PyMemberDef pywinpcap_capture_members[] = {
    {"device", T_OBJECT_EX, offsetof(pywinpcap_capture_object, m_device), 0, "device"},
    {NULL}  /* Sentinel */
};

static PyMethodDef pywinpcap_capture_methods[] = {
	{"next",	(PyCFunction)pywinpcap_capture_next,	METH_VARARGS,	pywinpcap_capture_next_doc},
	{"loop",	(PyCFunction)pywinpcap_capture_loop,	METH_VARARGS,	pywinpcap_capture_loop_doc},
	{"send",	(PyCFunction)pywinpcap_capture_send,	METH_VARARGS,	pywinpcap_capture_send_doc},
    {NULL}  /* Sentinel */
};

PyTypeObject pywinpcap_capture_type = {
    PyObject_HEAD_INIT(NULL)
    0,                         /*ob_size*/
    "pywinpcap.Capture",        /*tp_name*/
    sizeof(pywinpcap_capture_object), /*tp_basicsize*/
    0,                         /*tp_itemsize*/
	(destructor)pywinpcap_capture_dealloc,  /*tp_dealloc*/
    0,                         /*tp_print*/
    0,                         /*tp_getattr*/
    0,                         /*tp_setattr*/
    0,                         /*tp_compare*/
    (reprfunc)pywinpcap_capture_str, /*tp_repr*/
    0,                         /*tp_as_number*/
    0,                         /*tp_as_sequence*/
    0,                         /*tp_as_mapping*/
    0,                         /*tp_hash */
    0,                         /*tp_call*/
    (reprfunc)pywinpcap_capture_str,  /*tp_str*/	
    0,                         /*tp_getattro*/
    0,                         /*tp_setattro*/
    0,                         /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT,        /*tp_flags*/
    "PCap Reader object",      /* tp_doc */
    0,		               /* tp_traverse */
    0,		               /* tp_clear */
    0,		               /* tp_richcompare */
    0,		               /* tp_weaklistoffset */
    0,		               /* tp_iter */
    0,		               /* tp_iternext */
	pywinpcap_capture_methods,  /* tp_methods */
    pywinpcap_capture_members,  /* tp_members */
    0,                         /* tp_getset */
    0,                         /* tp_base */
    0,                         /* tp_dict */
    0,                         /* tp_descr_get */
    0,                         /* tp_descr_set */
    0,                         /* tp_dictoffset */
	(initproc)pywinpcap_capture_init,	/* tp_init */
    0,                         /* tp_alloc */
	pywinpcap_capture_new,      /* tp_new */
};
